import os, subprocess, sys
succ_seeds = {
    # "vi" : [0, 1, 6, 10],
    # "vi" : [6], 
    # "tr" : [1, 2, 4, 5, 6, 8, -1],
}
SDIorMUSE = 'SDI'

def vecmap_raw_step_iter():
    iter_rounds = 10
    src = sys.argv[1]
    trg = sys.argv[2]
    seed = 13
    norm_ops = ['center']
    for i in range(1, 2):
        exp_name = 'Nov17_iternormcent'
        exp_id = src + '-' + trg + '_vecmap_raw'
        norm_ops += ['center'] if i%2==0 else ['unit']
        subprocess.run([
            "python3", "vecmap_main.py",
            "--unsupervised",
            "data/wiki." + src + ".vec",
            "data/wiki." + trg + ".vec",
            f"dumped/{exp_name}/{exp_id}/vectors-{src}.txt",
            f"dumped/{exp_name}/{exp_id}/vectors-{trg}.txt",
            "--validation", f"data/crosslingual/dictionaries/{src}-{trg}.0-5000.txt",
            "--cuda" ,
            "--seed", str(seed),
            "--verbose",
            "--log", f"dumped/{exp_name}/{exp_id}/training_log.txt",
            "--normalize"]
            + norm_ops
        )

        with open(f"dumped/{exp_name}/{exp_id}/training_log.txt", 'a') as logfile:
            subprocess.run([
                "python3","eval_translation.py",
                f"dumped/{exp_name}/{exp_id}/vectors-{src}.txt",
                f"dumped/{exp_name}/{exp_id}/vectors-{trg}.txt",
                "-d", f"data/crosslingual/dictionaries/{src}-{trg}.5000-6500.txt",
                "--cuda",
                "--seed", str(seed),
            ], stdout=logfile)

def vecmap_raw_with_random_seeds():
    iter_rounds = 10
    src = sys.argv[1]
    trg = sys.argv[2]
    norm_ops = ['unit', 'center'] * 5 + ['unit']
    for seed in range(0, 1):
        exp_name = 'Nov17_iternormcent'
        exp_id = src + '-' + trg + '_vecmap_raw'
        exp_path =  f"dumped/{exp_name}/{exp_id}/"
        os.makedirs(exp_path, exist_ok=True)
        subprocess.run([
            "python3", "vecmap_main.py",
            "--unsupervised",
            "--orthogonal",
            f"data/wiki.{src}.vec",
            f"data/wiki.{trg}.vec",
            exp_path + f"vectors-{src}.txt",
            exp_path + f"vectors-{trg}.txt",
            "--validation", f"data/crosslingual/dictionaries/{src}-{trg}.0-5000.txt",
            "--cuda" ,
            "--seed", str(seed),
            "--verbose",
            "--log", f"dumped/{exp_name}/{exp_id}/training_log.txt",
            "--normalize"]
            + norm_ops
        )

        log_file_path = exp_path + "training_log.txt"
        file_mode = 'a' if os.path.exists(log_file_path) else 'w'
        with open(log_file_path, file_mode) as logfile:
            subprocess.run([
                "python3","eval_translation.py",
                f"dumped/{exp_name}/{exp_id}/vectors-{src}.txt",
                f"dumped/{exp_name}/{exp_id}/vectors-{trg}.txt",
                "-d", f"data/crosslingual/dictionaries/{src}-{trg}.5000-6500.txt",
                "--cuda",
                "--seed", str(seed),
                "--retrieval", "csls",
            ], stdout=logfile)

def muse_with_normalizestored():
    src = sys.argv[1]
    trg = sys.argv[2]
    seeds = succ_seeds[src] if src in succ_seeds else range(0, 20)
    # seeds = [0]
    exp_name = 'Nov17_iternormcent'
    exp_id = src + '-' + trg
    norm_ops = 'center,renorm,' * 5
    for i in seeds:
        subprocess.run([
            "python3", "muse_main.py", 
            "--src_lang", src,
            "--tgt_lang", trg,
            # "--src_emb", "dumped/tmp_itrecent_wordvec/vectors-" + src + ".pth",
            # "--tgt_emb", "dumped/tmp_itrecent_wordvec/vectors-" + trg + ".pth",
            "--src_emb", f"data/wiki.{src}.vec",
            "--tgt_emb", f"data/wiki.{trg}.vec",
            "--n_refinement", "5",
            # "--n_adversarials", "0",
            # "--use_sdi", "--seeddict_sdi",
            # "--sdi_cutoff", "20000",
            "--max_vocab", "200000",
            # "--n_epochs", "5",
            # "--epoch_size", "1000000",
            # "--map_optimizer", "sgd,lr=0.02",
            # "--dis_optimizer", "sgd,lr=0.02",
            # "--batch_size", "128",
            # "--exp_id", "jul29_renorm",
            "--exp_name", exp_name, 
            "--exp_id", exp_id, 
            "--normalize_embeddings", norm_ops,
            "--seed", str(i),
            # "--load_temp_model", "~/Documents/MI_Tlab/CUWTr/dumped/debug/jul25_19_32/"
            # "--dico_eval", "./datasets/crosslingual/dictionaries/en-es.txt"
            "--cuda", 
            "--export", "",
        ])
if __name__ == '__main__':
    vecmap_raw_with_random_seeds()
